#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：Flower_Class
@File    ：windows.py
@IDE     ：PyCharm
@Author  ：Rice_cc
@Date    ：2024/10/16 10:07
"""

import tensorflow as tf
from PyQt5.Qt import *
import sys
import cv2
from PIL import Image
import numpy as np

# import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"


# 读取图像，解决imread不能读取中文路径的问题
def cv_imread(filePath):
    """
    读取图像，解决imread不能读取中文路径的问题
    Args:
        filePath (str): 文件路径
    Returns:
        opencv_image: 输出读取的图像变量
    """
    cv_img = cv2.imdecode(np.fromfile(filePath, dtype=np.uint8), -1)
    # imdecode读取的是rgb，如果后续需要opencv处理的话，需要转换成bgr，转换后图片颜色会变化
    # cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
    return cv_img


# noinspection PyPackageRequirements
class MainWindow(QWidget):
    def __init__(self):     # 初始化界面
        super().__init__()  # 继承父类
        self.img_label = None
        self.result = None
        self.setWindowIcon(QIcon('ICON/LOGO.png'))
        self.setWindowTitle('花卉识别')     # 设置窗口标题
        self.model = tf.keras.models.load_model('models/mobilenet_flower.h5')
        self.to_predict_name = "ICON/init_Show.jpg"
        self.class_name = ['雏菊', '蒲公英', '玫瑰', '向日葵', '郁金香']
        self.resize(700, 500)   # 窗口大小
        self.initUI()           # 调用窗口控件函数

    # 生成界面控件
    def initUI(self):
        # main_widget = QWidget()
        main_layout = QHBoxLayout()     # 水平布局
        font = QFont('楷体', 15)         # 设置字体

        # 左边可视化控件
        left_widget = QWidget()
        left_layout = QVBoxLayout()     # 左边的窗口垂直布局
        img_title = QLabel("图像")
        img_title.setFont(font)         # 设置标签的字体格式
        img_title.setAlignment(Qt.AlignCenter)  # 设置标签居中
        self.img_label = QLabel()
        img_init = cv2.imread(self.to_predict_name)     # 读入数据
        img_init = cv2.resize(img_init, (224, 224)) # 设置读入的图像为224*224
        cv2.imwrite('ICON/target.png', img_init)
        # 将图像显示在标签img_label上
        self.img_label.setPixmap(QPixmap('ICON/target.png'))
        left_layout.addWidget(img_title)    # 添加标签控件 img_title
        left_layout.addWidget(self.img_label, 1, Qt.AlignCenter)    # 添加标签控件
        left_widget.setLayout(left_layout)      # 将左边的widget设置为垂直布局

        # 右边可视化控件
        right_widget = QWidget()                            # 右边添加窗口控件
        right_layout = QVBoxLayout()                        # 右边垂直布局
        btn_change = QPushButton("输入图像")                  # 添加按钮控件btn_change
        btn_change.clicked.connect(self.change_img)         # 建立信号和槽的连接
        btn_change.setFont(font)                            # 设置按钮字体
        btn_predict = QPushButton("识别")                    # 添加按钮控件btn_predict
        btn_predict.setFont(font)                           # 设置按钮字体
        btn_predict.clicked.connect(self.predict_img)       # 建立信号和槽的连接

        label_result = QLabel('识别结果')                     # 添加控件标签label_result
        label_result.setFont(QFont('楷体', 16))              # 设置标签字体格式
        self.result = QLabel("待识别")                        # 添加控件self.result
        self.result.setFont(QFont('楷体', 24))               # 设置标签字体格式
        right_layout.addStretch()                           # 添加弹簧
        # 右边布局中加入标签label_result, 居中显示
        right_layout.addWidget(label_result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(self.result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(btn_change)
        right_layout.addWidget(btn_predict)
        right_layout.addStretch()
        right_widget.setLayout(right_layout)                # 将右边的widget设置为垂直布局

        # 窗口布局
        main_layout.addWidget(left_widget)                  # 主窗口加入子控件窗口 left_widget
        main_layout.addWidget(right_widget)
        self.setLayout(main_layout)                         # 主窗口设置水平布局

    # 定义槽函数 change_img
    # noinspection PyPackageRequirements
    def change_img(self):
        # 选择文件
        openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Image files(*.jpg , *.png)')
        img_name = openfile_name[0]
        # noinspection PyPackageRequirements
        if img_name == '':
            pass
        else:
            self.to_predict_name = img_name
            # img_init = cv2.imread(self.to_predict_name)
            img_init = cv_imread(self.to_predict_name)
            img_init = cv2.resize(img_init, (224, 224))
            cv2.imwrite('ICON/target.png', img_init)
            self.img_label.setPixmap(QPixmap('ICON/target.png'))

    # 定义槽函数 predict_img
    def predict_img(self):
        img = Image.open('ICON/target.png')
        img = np.asarray(img)                   # 将输入的转换成数组
        outputs = self.model.predict(img.reshape(1, 224, 224, 3))
        result_index = np.argmax(outputs)       # 取出数组中最大值的索引
        result = self.class_name[result_index]  # 读出数组class_name对应索引的数据
        result_confidence = outputs[0][result_index]  # 获取最大概率值，即置信度
        self.result.setText(f'{result}的概率是{result_confidence:.2f}')             # 显示识别结果


if __name__ == '__main__':
    # gpu = tf.test.gpu_device_name()
    # print(gpu)

    app = QApplication(sys.argv)      # pyqt窗口需要在QApplication方法中使用
    x = MainWindow()                    # 生成MainWindow类的实例 x
    x.show()                            # x调用show方法
    sys.exit(app.exec_())               # 消息结束时，进程结束并返回0，再调用sys.exit(0)退出程序

